using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace JufGame
{
    public abstract class BaseLayer : IANNLayer
    {
		public ComputeShader compute;
		public ComputeBuffer outputBuffer, dInputBuffer;
		protected int kernelForward, kernelBackward;
        protected int threadGroupsX, threadGroupsY;
		protected int batch, inputCount;

		/// <summary>
		/// 初始化激活函数层
		/// </summary>
		/// <param name="shader">使用的激活函数compute shader</param>
		/// <param name="batch">输入数据的批数</param>
		/// <param name="inputCount">每组输入数据的长度</param>
        public BaseLayer(ComputeShader shader, int batch, int inputCount)
		{
			compute = shader;
			this.batch = batch;
			this.inputCount = inputCount;
            kernelForward = compute.FindKernel("Forward");
            kernelBackward = compute.FindKernel("Backward");
			threadGroupsX = Mathf.CeilToInt(inputCount / 8.0f);
			threadGroupsY = Mathf.CeilToInt(batch / 8.0f);
		}

		/// <summary>
		/// 前向传播，记得设置矩阵维度
		/// </summary>
		/// <param name="input">输入数据</param>
		public virtual void Forward(ComputeBuffer inputBuffer)
		{
			compute.SetBuffer(kernelForward, "inputBuffer", inputBuffer);
			compute.SetBuffer(kernelForward, "outputBuffer", outputBuffer);
			compute.Dispatch(kernelForward, threadGroupsX, threadGroupsY, 1);
		}

		/// <summary>
		/// 基于compute shader的反向传播，记得设置矩阵维度
		/// </summary>
		/// <param name="dout">求导数据</param>
		public virtual void Backward(ComputeBuffer dOutputBuffer)
		{
			dInputBuffer ??= new ComputeBuffer(batch * inputCount, sizeof(float));
			compute.SetBuffer(kernelBackward, "dInputBuffer", dInputBuffer);
			compute.SetBuffer(kernelBackward, "dOutputBuffer", dOutputBuffer);
			compute.Dispatch(kernelBackward, threadGroupsX, threadGroupsY, 1);
		}
	
		/// <summary>
		/// 释放 ComputeBuffer 资源
		/// </summary>
        public virtual void Release()
        {
            outputBuffer?.Release();
			dInputBuffer?.Release();
        }
    }
}
